'''magtorrent -- small BT implementation tailored for snarfing metadata.
Will deviate from the protocol where it suits us for performance.
'''
#  Copyright (c) 2010
#  The Regents of the University of Michigan
#  All Rights Reserved

#  Permission is granted to use, copy, create derivative works, and
#  redistribute this software and such derivative works for any purpose,
#  so long as the name of the University of Michigan is not used in
#  any advertising or publicity pertaining to the use or distribution
#  of this software without specific, written prior authorization. If
#  the above copyright notice or any other identification of the
#  University of Michigan is included in any copy of any portion of
#  this software, then the disclaimer below must also be included.

#  This software is provided as is, without representation or warranty
#  of any kind either express or implied, including without limitation
#  the implied warranties of merchantability, fitness for a particular
#  purpose, or noninfringement.  The Regents of the University of
#  Michigan shall not be liable for any damages, including special,
#  indirect, incidental, or consequential damages, with respect to any
#  claim arising out of or in connection with the use of the software,
#  even if it has been or is hereafter advised of the possibility of
#  such damages.

import event

import cStringIO
import logging
import os
import socket


UINT_STRUCT = struct.Struct('>I')

def generate_random_peer_id():
  '''Return a random BitTorrent peer id.'''
  return ''.join(chr(random.randint(0, 255)) for _ in xrange(20))


def supports_magnet_protocol(reserved_bytes):
  return bool(reserved_bytes[5] & 0x10)


class PeerConnection(object):
  '''A connection to a BitTorrent peer.

  Supports only the initial handshake and querying the peer for
  metadata. See BEPs 0003, 0009, and 0010.
  '''

  # Derived from Encrypter.py in Bram Cohen's MIT-licensed BT implementation.
  # Heavily stripped down.

  BT_HANDSHAKE = '\x13BitTorrent protocol'

  # Support only magnet protocol extension.
  RESERVED_BYTES = '\x00\x00\x00\x00\x00\x10\x00\x00'

  FIRST_MESSAGE_LEN = len(BT_HANDSHAKE) + len(RESERVED_BYTES)

  UT_METADATA_PIECE_SIZE = 16 * 1024

  def __init__(self, peer_address, torrent_hash=None):
    self.torrent_hash = torrent_hash
    self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    self.sock.setblocking(0)
    self.sock.connect(peer_address)
    self.peer_id = generate_random_peer_id()
    self.other_id = None

    self.next_len = PeerConnection.FIRST_MESSAGE_LEN
    self.next_func = self.read
    self.recv_buffer = cStringIO.StringIO()

    self.last_piece = None

    first_msg = PeerConnection.BT_HANDSHAKE + PeerConnection.RESERVED_BYTES
    if torrent_hash is not None:
      first_msg += torrent_hash
    self.sock.send(first_msg)

    event.read(self.sock, self._handle_recv)

  def close(self):
    '''Close the connection.'''
    # TODO: does creator need a callback for this?
    if not self.closed:
      self.sock.close()
      self.closed = True

  def _handle_recv(self):
    # This is an event handler; return value indicates whether to reregister.
    buf = self.recv_buffer
    recvd = self.sock.recv(4096)

    # Process the data, one protocol message at a time.
    # Invariant: self.next_len is the length of the next data item to read and
    # self.next_func is the method that reads and processes the next item.
    while True: 
      if self.closed:
        return False

      bytes_needed = self.next_len - buf.tell()

      if len(recvd) < bytes_needed:
        buf.write(recvd)
        return True

      buf.write(recvd[:bytes_needed])
      recvd = recvd[bytes_needed:]

      msg = buf.getvalue()
      buf.reset()
      buf.truncate()

      try:
        x = self.next_func(msg)
      except:
        self.next_len, self.next_func = 1, lambda x,y: None
        raise

      if x is None:
        self.close()
        return
      self.next_len, self.next_func = x

  def _read_first_message(self, data):
    if not data.startswith(PeerConnection.BT_HANDSHAKE):
      return None
    data = data[len(PeerConnection.BT_HANDSHAKE):]
    if not supports_magnet_protocol(data):
      return None

    return 20, self._read_torrent_hash

  def _read_torrent_hash(self, data):
    if self.torrent_hash:
      if data != self.torrent_hash:
        return None
    else:
      # Gotta trust the peer re: hash. TODO: don't support this.
      self.torrent_hash = data
      self.sock.send(data)

    self._send_extension_handshake()
    return 20, self._read_peer_id

  def _read_peer_id(self, data):
    self.other_id = data
    self.sock.send(self.peer_id)

    return 4, self._read_msg_len

  def _read_msg_len(self, data):
    l = UINT_STRUCT.unpack(data)[0]
    return l, self._read_msg

  def _read_msg(self, data):
    if ord(msg[0]) == 20:  # Extension message.
      self._handle_extension_message(msg[1:])
    return 4, self._read_msg_len

  def _handle_extension_message(self, msg):
    extended_msg_id = ord(msg[0])
    if extended_msg_id == 0:
      body = BitTorrent.bencode.bdecode(msg[1:])
      if 'ut_metadata' not in body['m']:
        logging.error('Client (%s) set magnet bit but didn\'t body it'
                      % body.get(v, ''))
        return None
      self.extension_messages = body['m']
      self.metadata_size = body['metadata_size']
      self.got_magnet_handshake()
    elif extended_msg_id == self.extension_messages['ut_metadata']:
      # May be either a bencoded dict or a bencoded dict with crap attached.
      msg_type = body.get('msg_type', -1)
      if msg_type == 0: # Ignore metadata request message.
        pass
      elif msg_type == 1: # Data.
        

  def _send_extension_handshake(self):
    pass # Looks like we don't actually have to do this.
